package edu.zju.gis.dldsj.tasks;

import edu.zju.gis.dldsj.config.CommonSetting;
import edu.zju.gis.dldsj.entity.*;
import edu.zju.gis.dldsj.service.WfService;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;

import java.util.Comparator;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 对工作流任务的监控
 *
 * @author yanlong_lee@qq.com
 * @version 1.0 2018/11/06
 */
@Slf4j
@AllArgsConstructor
public class WorkflowMonitor implements Runnable {
    private CommonSetting setting;
    private WfService wfService;
    private WfDag wfDag;
    private List<WfNode> wfNodes;
    private String workspace;
    /**
     * 最终结果的存储位置
     */
    private String resultAddress;

    @Override
    public void run() {
        log.info("开启新的工作流监控任务：dagId=`{}` and userId=`{}`", wfDag.getDagId(), wfDag.getUserId());
        try {
            AirflowDagRun dagRun = null;
            //1. 轮询查找dag运行实例
            for (int loop = setting.getJobFailInterval() / setting.getMonitorInterval(); loop >= 0; loop--) {
                List<AirflowDagRun> dagRuns = wfService.getDagRun(wfDag.getDagId());
                if (!dagRuns.isEmpty()) {
                    dagRuns.sort(Comparator.comparingInt(AirflowDagRun::getId).reversed());
                    dagRun = dagRuns.get(0);
                    break;
                }
                Thread.sleep(setting.getMonitorInterval());
            }
            if (dagRun == null) {
                log.error("没有DAG`{}`的实例正在运行", wfDag.getDagId());
            } else {
                //2. 记录dag运行实例到业务数据库
                WfRun wfRun = new WfRun();
                wfRun.setId(dagRun.getId());
                wfRun.setDagId(dagRun.getDagId());
                wfRun.setExecutionDate(dagRun.getExecutionDate());
                wfRun.setRunId(dagRun.getRunId());
                wfRun.setState(dagRun.getState());
                wfRun.setWorkspace(workspace);
                wfRun.setResult(resultAddress);
                wfService.insert(wfRun);

                //3. 记录dag各子任务运行实例到业务数据库
                List<WfInstance> tasks = wfNodes.stream().map(node -> {
                    WfInstance instance = new WfInstance();
                    instance.setDagId(wfRun.getDagId());
                    instance.setTaskId(node.getTaskId());
                    instance.setExecutionDate(wfRun.getExecutionDate());
                    instance.setArtifactId(node.getArtifactId());
                    instance.setParams(node.getParams());
                    wfService.insert(instance);
                    return instance;
                }).collect(Collectors.toList());
                //4. 更新dag实例运行状态到业务数据库
                while (true) {
                    for (WfInstance task : tasks) {
                        AirflowTaskInstance instance = wfService.getTaskInstance(task.getDagId(), task.getTaskId(), task.getExecutionDate());
                        if (instance != null) {
                            task.setStartDate(instance.getStartDate());
                            task.setEndDate(instance.getEndDate());
                            task.setDuration(instance.getDuration());
                            task.setState(instance.getState());
                            wfService.update(task);
                        }
                    }
                    dagRun = wfService.getDagRun(dagRun.getId());
                    if (!DagState.RUNNING.getState().equals(dagRun.getState())) {
                        wfRun.setState(dagRun.getState());
                        wfService.update(wfRun);
                        break;
                    }
                    Thread.sleep(setting.getMonitorInterval());
                }
                log.info("工作流`{}`@`{}`监控已完成，最终状态为：{}", wfDag.getDagId(), wfRun.getExecutionDate(), wfRun.getState());
            }
        } catch (Exception e) {
            log.error("工作流`" + wfDag.getDagId() + "`监控失败", e);
        }
    }
}
